{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9902a6a3",
   "metadata": {},
   "source": [
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/langchain-ai/langchain-academy/blob/main/module-3/time-travel.ipynb) [![Open in LangChain Academy](https://cdn.prod.website-files.com/65b8cd72835ceeacd4449a53/66e9eba12c7b7688aa3dbb5e_LCA-badge-green.svg)](https://academy.langchain.com/courses/take/intro-to-langgraph/lessons/58239536-lesson-5-time-travel)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba98beac-d461-4d7d-878a-11beca03ea1c",
   "metadata": {},
   "source": [
    "# Time travel\n",
    "\n",
    "## Review\n",
    "\n",
    "We discussed motivations for human-in-the-loop:\n",
    "\n",
    "(1) `Approval` - We can interrupt our agent, surface state to a user, and allow the user to accept an action\n",
    "\n",
    "(2) `Debugging` - We can rewind the graph to reproduce or avoid issues\n",
    "\n",
    "(3) `Editing` - You can modify the state \n",
    "\n",
    "We showed how breakpoints can stop the graph at specific nodes or allow the graph to dynamically interrupt itself.\n",
    "\n",
    "Then we showed how to proceed with human approval or directly edit the graph state with human feedback.\n",
    "\n",
    "## Goals\n",
    "\n",
    "Now, let's show how LangGraph [supports debugging](https://langchain-ai.github.io/langgraph/how-tos/human_in_the_loop/time-travel/) by viewing, re-playing, and even forking from past states. \n",
    "\n",
    "We call this `time travel`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd48aeb6-8478-4cb4-aef1-d524b80824d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install --quiet -U langgraph langchain_openai langgraph_sdk langgraph-prebuilt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d32093f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, getpass\n",
    "\n",
    "def _set_env(var: str):\n",
    "    if not os.environ.get(var):\n",
    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
    "\n",
    "_set_env(\"OPENAI_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0497d316-832a-4668-b133-fd317ee81220",
   "metadata": {},
   "source": [
    "Let's build our agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d64ab3a1-b39c-4176-88c7-791a0b80c725",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "def multiply(a: int, b: int) -> int:\n",
    "    \"\"\"Multiply a and b.\n",
    "\n",
    "    Args:\n",
    "        a: first int\n",
    "        b: second int\n",
    "    \"\"\"\n",
    "    return a * b\n",
    "\n",
    "# This will be a tool\n",
    "def add(a: int, b: int) -> int:\n",
    "    \"\"\"Adds a and b.\n",
    "\n",
    "    Args:\n",
    "        a: first int\n",
    "        b: second int\n",
    "    \"\"\"\n",
    "    return a + b\n",
    "\n",
    "def divide(a: int, b: int) -> float:\n",
    "    \"\"\"Divide a by b.\n",
    "\n",
    "    Args:\n",
    "        a: first int\n",
    "        b: second int\n",
    "    \"\"\"\n",
    "    return a / b\n",
    "\n",
    "tools = [add, multiply, divide]\n",
    "llm = ChatOpenAI(model=\"gpt-4o\")\n",
    "llm_with_tools = llm.bind_tools(tools)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1d8622a9-57cd-44dc-8696-46c5ab32d0b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "/9j/4AAQSkZJRgABAQAAAQABAAD/4gHYSUNDX1BST0ZJTEUAAQEAAAHIAAAAAAQwAABtbnRyUkdCIFhZWiAH4AABAAEAAAAAAABhY3NwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQAA9tYAAQAAAADTLQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAlkZXNjAAAA8AAAACRyWFlaAAABFAAAABRnWFlaAAABKAAAABRiWFlaAAABPAAAABR3dHB0AAABUAAAABRyVFJDAAABZAAAAChnVFJDAAABZAAAAChiVFJDAAABZAAAAChjcHJ0AAABjAAAADxtbHVjAAAAAAAAAAEAAAAMZW5VUwAAAAgAAAAcAHMAUgBHAEJYWVogAAAAAAAAb6IAADj1AAADkFhZWiAAAAAAAABimQAAt4UAABjaWFlaIAAAAAAAACSgAAAPhAAAts9YWVogAAAAAAAA9tYAAQAAAADTLXBhcmEAAAAAAAQAAAACZmYAAPKnAAANWQAAE9AAAApbAAAAAAAAAABtbHVjAAAAAAAAAAEAAAAMZW5VUwAAACAAAAAcAEcAbwBvAGcAbABlACAASQBuAGMALgAgADIAMAAxADb/2wBDAAMCAgMCAgMDAwMEAwMEBQgFBQQEBQoHBwYIDAoMDAsKCwsNDhIQDQ4RDgsLEBYQERMUFRUVDA8XGBYUGBIUFRT/2wBDAQMEBAUEBQkFBQkUDQsNFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBT/wAARCADbAMcDASIAAhEBAxEB/8QAHQABAAICAwEBAAAAAAAAAAAAAAUGBwgBAwQJAv/EAFcQAAEDBAADAgcIDAkJCQEAAAECAwQABQYRBxIhEzEIFBYiQVGUFRcjVVZh0dMyNkJUcXSBkZOVtNIJNThSU3WSstQkM2JjcnOhs8EYNEVXgoOEscPx/8QAGgEBAQADAQEAAAAAAAAAAAAAAAECAwQFB//EADURAQABAwAFCgMIAwAAAAAAAAABAgMRBBIhMVETFEFSYXGRobHBFSPRBSIyM1OB4fBCQ8L/2gAMAwEAAhEDEQA/APqnSlKBSlKBSvJdLnHs1vfmylFLDKeY8qSpSj3BKUjqpROgEjqSQB1NQfk9Lyb4e/OOsxVbLdnjulCEJ9HbKSduL9YB5BvQCtc6ttNETGtVOI/u5cJmTfbbCcKJFwisLHQpdfSkj8hNdPlVZfjiB7Sj6a6o+F4/EbCGLFbWkAAaREbHd0Hort8lbL8TwPZkfRWfye3yNh5VWX44ge0o+mnlVZfjiB7Sj6aeStl+J4HsyPop5K2X4ngezI+inye3yXYeVVl+OIHtKPpp5VWX44ge0o+mnkrZfieB7Mj6KeStl+J4HsyPop8nt8jYeVVl+OIHtKPprlGTWdxQSi7QVKPoTJQT/wDdceStl+J4HsyPorheJ2NxBSqzW9ST0IMVBB/4U+T2+RsSiVBaQpJCkkbBB2CK5qsLwKDBWp+wKVjssnm/yIajrP8ArGPsFA+kgBXfpQJ3UjY7y5PW/DmMeKXOLoPMg7QsHucbPpQrR0e8EEHqKxqojGtROY8JTHBLUpStKFKUoFKUoFKUoFKUoFKUoFKUoKvdtXbOLTbl6VGgsLuTiD907zBtn8IG3VdfSEHvGxaKrDo8T4ksOL2ET7WppCtdOZl3m1v1kPEj/ZPqqz10Xd1ERux9c+ayUpSudFAhceMHuWUXLHYd4cmXa3KfRIajQJLiA4ykqdbS6lsoW4kA7QlRVsa1vpVZ4U+E9jfEPhnMzC4NS7AxAK1TUPwJXZtI7dxprkcUykPKIQNhvmKSrRAPSqjhwvGOeEAYOF2TLbZityudwkZNBvluKLU25yqUmZCkK9LroSezQpQIWSUoIquYvc86w7wd7hhFnx3J7VllinuplzI1rUrtITlzUp12A4oFt93xdwqSkbOwemwKDOVq8ILAbziGQZPFv27Rj6Su6qdhyGn4aeXm2thbYdGx1HmddHW9VVM78LHFMYtNjuNrbn3yHcb3GtSpLNrm9kG3DtbzSgwQ/pPVIbJ5yfNJ1qsG3bDbxLsvH1NmxvO5MPIcQiItb2RsSpEue8yZCXEjtOZxKtup5WlBKtbKU8vWs7cfrDcU8PcHm2myzLonGshtN1k262sFyT4swsBwNNDqtSQd8o69DQZfs92j320w7lE7bxWWyl9rxhhbDnKobHM24ErQdHqlQBHcQK9lRuOXxvJbJEubUSbAbkp50x7lGXGkIGyNLbWApJ6b0R6RUlQKrGXatdzsN5RpK25iIDx6+ezIUGwn9KWVfkPrqz1WM8T43Fs9vSCXZd1iFIA30ZdEhRPqHKyrr84rosfmRE7unu6fJY3rPSlK50KUpQKUpQKUpQKUpQKUpQKUpQRWRWZV4iNFhxLFwiOiTDfWCQ26AR1AIJSpKlIUAeqVqAI7667XfI18D9vlNCNcUJKZNueOzy9xUnYHO2d9FgaPcdEFImajrzj1uyFptu4RG5PZEqacO0uNKI0VIWNKQddNpINbqaqZjVr3en9/vbe9SB4NnCdJBHDfFgR3EWhj92uP+zXwn/8ALbFf1Qx+7VhODFvpHyK+x0dAEeOB3Q/C4lSj+U7p5EyPlVfv0zP1VZalvr+UmI4rJHjtRI7TDLaWmWkhCG0DSUpA0AB6ABXZVX8iZHyqv36Zn6qnkTI+VV+/TM/VU5O31/KTEcVopWvvgtXrIeMfBe05VfsouqLnKky2nBDU023ytSXG06BbJ+xQN9e+steRMj5VX79Mz9VTk7fX8pMRxeDIuB3DzLrzIu17wiwXe6SeXtpk23NOuucqQlPMpSSTpKQPwAVHq8G/hStKArhxi6ggcqQbSweUbJ0PN9ZJ/LU/5EyPlVfv0zP1VBhLxBCsnvy0nprt2h/xDYNOTt9fykxHF3Wy04vwtx0RbdCt2NWZtZUmPEaSw12ij3JQkDalH0AbJ7tmubPCkXW7C+z2DGKWlMwYq/s2m1EFS1j0LVyp6fcgAd5VXZa8LtVqmiaGnZlwAIEyc+uQ6nfeEqWTyA+pOh81TtSaqaImLfT0/Q2RuKUpWhClKUClKUClKUClKUClKUClKUClKUClKUClKUGu/gB/yYce/Hbj+2vVsRWu/gB/yYce/Hbj+2vVsRQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQKUpQa7+AH/Jhx78duP7a9WxFa7+AH/Jhx78duP7a9WxFApSlApSlApSlApSlApSlApSlApSlApSlApSlApSq3e8oksXFdttENqbNaQlchyQ8WmWAr7EEhKipZAJ5QOgAJKeZO9lFuq5OKV3rJSqR7u5h94WP2t76unu7mH3hY/a3vq66Oa18Y8YMLvWjf8J/wLVlOE2ziTbI5cuNgAh3HkGyqEtZKFf8AtuKPd6HVE9E1tb7u5h94WP2t76uo/IU5FlVhuNlulnsMu23CO5Ekx1y3tONLSUqSfg/SCac1r4x4wYfMf+Dy4KOcU+O8K9yW1CyYkpu6vuDYCpIVuM3sdx508/qIaUPTX1/rXXwdODF08HDBXccszFpuCpEtyZJnyJDiXHlK0EggN6ASgJTodN7PTmNZT93cw+8LH7W99XTmtfGPGDC70qke7uYfeFj9re+rp7u5h94WP2t76unNa+MeMGF3pVLRmF4tI7e+W6Ei3D/OyYElbimB/PUhSBtA6bIOx360CaulaLlqq3+IxgpSlakKUpQKUpQKUpQKUpQKUpQKUpQKoVoPNleZ79FyaA6ejxGMf+pq+1QbP9teaf1m1+wxa7tF/wBnd/1DKN0pulYE4/Z1kdsyI2vDcjvLN5h2pVwetFmskWYlKeZYQ7JekKSENqKSkIQQs8qiN1W8j4wZnPs/D3KZN7fwbCbxjkefNvVutDdwZYuLvIezkhYUpqPyq6LGupPMsVlNUQxbP15LfeIF3MoQZsaaYj6o0gR3Uudi8kAqbXonlUNjaT1GxWEPLrKIPHpy25JksjHMdmS2mseiItTTtuvDSmQSgzNFSJBXz6QVJ2EjlCt1SIWT8QrRh2R+5NxlSEQM9mwr1e7RY4jtxRCQ0nTwjIaSh1fPyBauRS+XuB10aw2vU4hCkJUpKVLOkgnqTrfSv1WrmQtXDPeKXAy42XiRLktTbRdi1ebdAiJDpQlkrcDbjSwlSwQhSSPN7PoEnm3xxv40ZVh+QZLeMUv11u9sxqRGbuFrYscU2uOT2XaMPS1rS8pwpXzfA75OdIUO801htJX5LiA4lBUkLUCQknqQNbOvyj89a/ZZkue3jOOLUOzZibBBxK3xJsGO3bY7/auuRVuFDi3Ek9mS31A0rzuigBqoK1u3viRx44ZZGxks2wLuuAG6uRoceM4hIU9DW4yO1bWeVZWNnfMOQcpGztrDYbNQDht+BAI8QkdCNg/Bqq5WklVqhEkkllBJP+yKp2a/abfvxB//AJaquFn/AIphf7hH90Vb/wCTT3z6QvQ9lKUrzkKUpQKUpQKUpQKUpQKUpQKUpQKoNn+2vNP6za/YYtX6qPdYsrG8guM9MKROt9zWh5a4bRccYdS2hohSAOYpKUJIUN6IIIHTfbosxmqnpmPeJ9mUdKo5jwStWYZS9fjeL5ZZMyEi3XFm0TAw3cI6VKKUO+aVDXaLAU2pCtKI3UFP8Ge0T8QteLeVmWR8ehW4Wly3R7g2hqZFCiQ28Oy/mq5OZHKopABJ76u134lWewWyVcrozdrdb4rZdfly7TKaaaQO9SlqbASB6ya67FxSseUWmNdbMi6XW2SU87EyFapLzLo2RtK0tkEbBHQ+iurkK+rJqzwQU7gLabpllvvEy+5BKgW+axcImPuzUm3MSGUBLS0I5OcBOgQnn5ebqQTXL/A2GiFc2bXlOS4+7cL2/fnpNrmNNuds6kJW3otKSprQ2EqSog9d7A1a/LON8WX79SS/qqeWcYf+GX79SS/qqchX1ZNWeCkveDjjjeN4larVcr3YHsXU+q3XS2y0iWC/vxjnUtC0r7QqKlbT392q8WTeC/juVIyGNJv2SxrTf1iRcbVEnpbjvyeRKPGTpvm5/g0KI5uQqSCUHuq5Y/xXx/LLW3c7Gbjeba6VJRMt9skvsrKVFKgFobIJCgQevQgipHyzjfFl+/Ukv6qpyFfVNWeCJi8KLZHuGYTlzrhIlZTDjwp7jq2+iWmVMpUgBAAUUrJO9jfcAOlQk7wf7M/Bw1uBfL9Yp2K20WmFc7ZJbbkOxeRtBbe5m1IUD2SFHSRojY1Vx8s43xZfv1JL+qqGyfjPi2FR4z+QyZtiYkuhhl25W6RHS64e5CStABUfUOtXkK+rJqzwT+a/abfvxB//AJaquFn/AIphf7hH90VRJ8qTmFvk2m3224sGa0phyXOhLjNx0KBSpfwqQVEDekgHZI3pO1DIjLKY7LbSBpCEhKR6gOgrRpH3bdNE78z7E7IfulKV57EpSlApSlApSlApSlApSlApSvytaW0KUpQSlI2VE6AFB+qpXFjilD4SY5Husq0Xi+uSpjUCNAskNUmQ685vlGh0A6HqSPQBskA+e6cQr0c9xK02DFXr/jN3jOTJmUsS2xEiNBPwYT1JcUtRQRrXmnaebSuXv4VcK4XCe0XKFFu94vjtxnu3GTNvUxUl5bi9DQJ6ABKUjoOutnZoOiBgN7f4i5JfLzlT94xW5wUQYmJPRGxFjp0O1U5sbcUo8w6681ZB5tJ1eWGG4zLbLLaWmm0hCG0JCUpSBoAAdwFdlKBWuPh4cc/eW4Gz2oEjscjyLmtlv5TpbaVD4Z4ekcqDoEdyloNbHVgbwg/A4xDwk8ktt5ye+ZJDct8TxRiJa5TDbCRzqWpfK4ys86uYAkEbCE9OlBqb/Bd8dfcfI7pwvukgJiXTmuNqLivsZKUjtWh/ttpCgO4dkr0qr6V18+vAn8DDDMnxDDeKki8ZFHyKDdnZLceNKYTFUY0taUJUkslZSoNgKHON7VrW6+gtArplQ2JzXZSWG5DfMlfI6gKHMkhSTo+kEAg+ggV3UoMdysRv+IZNmWZWu93nJxOt/NFwyS+0mKmW2gBPYOKA7ILCUpI3ralKPMdamcIzsZPjdhm3i2P4jebq2tSbDd3EJlpWgnnSEg+eBrm2OvKUkhO9C11Vcy4W4rxAuuP3PILLHuVxsEtM22SnNhyM6Ck7SQRsEpSSk7B5RsdBQWqlYqdy3KuFTHEHI+Ic+33DCoLqZlods8J0zWY6iQpp5sbCuTzNKG97UokDonIeN5Fbsux+3Xu0SPG7XcY6JUWQEKR2jS0hSVaUARsEHqBQSVKUoFKUoFKUoFKUoFKUoIjKsusuD2R28ZBc41ntbS223JktwIbQpa0oRtR6DalJG/nqmXDH8k4k3POcXzOyW2Pw6lxkQoC4c53x6ZzJ26tZTyhtPUJCehBQfskkGrHxOsdoyLAL7Cv1lTkdq8WU+9alDfjXZfCpQOo6lSE6+fVccMMyb4g8PbBkbVtkWdu4xEPiBLSQ7H2NFCtgdxGt6699BK4zjNrw3H7fY7JCat1pgMpYjRWRpLaB3Aek/hPUnqak6UoFKUoFdMuWxAivSpTzcaMwhTjrzywlDaANlSiegAAJJNR+VZXZ8Hx6dfb9cWLVaILZdkS5KuVCE/8AUk6AA6kkAAk1rExByfw3Z7cq5Nz8R4EsuBbEAkszsnIOwtzXVuNsAgDqrvGzooC0+AAoL8F7HFpIUlUy4lKgdgjx17qK2KrxWWyW/G7TEtdqhMW62xG0sx4sZsIbaQBoJSkdAK9tApSlApSlBwRsaPUVTL5w1N1z7GcniZHeLQmzNOR3LRDkagzWVJOkutEa2lXKQoddJ16iLpSgpPDvNr9kaLs1lWKuYdNi3F2JFQ/MafbnMjzm3WlJOztJTsEdDsddEC7Vini/Fwl/P+Fq8omTY16avDirA3FBLb0nsjzJd0k6Ty+sjr6aytQKUpQKUpQKUpQKUr8rcQ2NrUEj/SOqDEnhDeEvj3g1Wyz3HJbLf7lAubrjCJNmitutsuJCVBDqnHEBKlgqKQNkhtf82tK4X8JznlymLsOPY5bbtdJ19W3bLheUlO4biylhhcdlSdOjaNrDqh3jR6Kr6D8SMExzivhd0xbI2WptquDRbcSVDmbV9y4gn7FaTog+givmhwj8FC88MfDixTGL0343ZYMpd6h3ZKfgpMdhKnGl9/RXaJbSpJO0k+kEE3Ej6s0rq8aZ/pm/7Qp40z/TN/2hTEjtqn8VeLGM8GMOl5LlVxTAt7HmoQPOdkOEea00jvWs67vwkkAEiC448fce4GY0xOuCXbteLg54taLFbxzyrjI6ANtpG9Dak7VrpsdCSlJxzwq4CZFnuYROKXGtTM/JmvPsmLNnmgWBBOx5vULf7tqO9EA7JCSmCJxXhdlPhS5DBzji3AcsuDxHBIx/h84o/CfzZM8fdKIPRs929EAcwXtO22hltDbaEobQAlKUjQAHcAK/VKBSlKBSlKBSlfhbqG9c60p33cx1QfuvJdn5cW1TXrfFROntsrXHiuvdil5wJJSgr5VcgJ0ObR1vej3V3eNM/wBM3/aFPGmf6Zv+0KuJHzoyD+FJQ5eIouXBeL4/apCykTrwFvRnRtKuQmKC2vvBPf6K298F3j3J8I7hs9lz+MLxVr3QdhsR1zPGg+hCEEupX2bfTmUtGtHq2evoGjfh0eC3Pe8I2xTcTjpci5/KDZCB8HHn7AeUsgealSSHST/rT3Jr6M8NcKs3C3ArFidnU2i32mKiM2dgFwjqpxWvulqKlH51GmJFqpXV40z/AEzf9oVyJDSiAHUEnuAUKYkdlKUqBSlKDy3Sb7m2yXL5ebsGVu8vr5Uk/wDSseWvErVfrdEuV5t8S8XKUyh56TOYS8ragCUp5h5qB3BI0ND17NXnKvtYvH4m9/cNV7GvtctX4o1/cFelo8zRbmqmcTlluh4ve+xb5NWf2Br92nvfYt8mrP7A1+7VF4V+EVYuJIykuNSbMixzJiFvTYclljxVhYT2y3nWkIQo75i0TzoG9joTVgwjjbhXEWe/CsN7EqW1H8bLL8Z6MpbG9ds32qE9o3sgc6Np6jr1FbYv3J/znxTM8U1732LfJqz+wNfu0977Fvk1Z/YGv3agMS48YJnV/RZrJkDc2e6lxcdJjvNNykt/Zlh1aAh4J9JbUrp17qrWD+EPa18HsTy7Npce1zr4XG241uivvF1xK3BpplAccOko2e/XedU5xc68+JmeLIZ4fYz0Ldgt0dwdUvRoyGXEH1pWgBST84IIqxYJdJF0sBMp0yJEaTIhqeOtuBp1SEqOgBzFKQToAb3rpXgsl5h5HZ4V1tz3jECayiQw9ylPO2obSrSgCNgjvFfrhn/Elw/rad+0LrC9VNyzM1TnEx7rnMbVupSleWxKUpQK8t0ukWy2+ROmvJjxGEFbjiu4AfMOpPqA6k9BXqrEHHW8uOzrNY0K0wUrnSE7+yKSEtD5xsrV+FCa7ND0edKv02uPosK5lXEW85Y+4lmRIs9q2Q3Fjr7N5xPoLjifOBP81JAG9Hm1uqaqw21xaluQI7ritcy3WgtSvwk9TXupX0ezao0enUtRiGOtKP8AJ61fFkP2dH0U8nrV8WQ/Z0fRUhVQvPFzEsfvLlrn3hDEppSUPHsXFNMKVrlS66lJQ2TsdFKHeK2VXYojNVWP3MzxT/k9aviyH7Oj6KeT1q+LIfs6Poqu3zjDiOOXOdb7hdizLgKQJaERXnBHCkJWlTikoISgpWnzyQnvG9ggevKOJmNYc/DZut0Sy/LQXWWmWnH1qbHe5ytpUQj/AEjofPWPL0Rn7+7ftMzxS/k9aviyH7Oj6KHHbUQR7mQ9Hp/3dH0VBcJ8ul55w7sl/nNsNSpzJccRGSUtg8yh5oJJ7gO8mrbWVFzXpiqJ2SZni77Jcbhi7iV2ae/bwkj4BKiphQ9RaPm/lAB9RFZx4fZ8zmcNbbyExbtHA8YjJO0kHoHEE96Tr8IPQ+gnA9eux3hzG8ltN1bVyhqQhl7r9kw4oIcB9ethWvWgV5Wn6DRpVuaoj78bp9pWJzsls3SlK+eiLyr7WLx+Jvf3DVexr7XLV+KNf3BVkyNlcjHro02kqcXFdSlI9JKCBVaxdaXMatKknaVRGSD6xyCvQs/kz3+y9DWa6YnkV44fcauGrWP3di93e73S7W6YuItNvmMuupebQJP2AUsbbKSQQd70KkMut978IHKbT7iYxfMPjWrG7zCkTL7BVB5X5kZLLUdoHq4EKHOVJBQOROiSa2cpTVRrDjyL3m7vBbHI+FXzGZGGSGZV4m3KCY8aOliG5HUww6fNeDiljRbJHKNnVQ2P2BVp4H4fa79jWdWfK8VuMyNDuuOWtUiRDf2s9shI5g9HdQ6Ek8qkq6g61sbb0pqio8JLjk124a47MzKImDk70RCp7CUhPK586QSEqI0SkdxJHoqx8M/4kuH9bTv2hdeuvNw1QU2GYv7ly6TlJOu8eMuDf/A//wArKvZYq74916FspSleahSlKBWEON0VUfNbVKV/m5UBbKTr7ptzmI/M6PzH1Vm+qzxAw5OaWExULSzOYWH4jy96Q4ARpWvuVAlJ+Y77wK9L7P0inRtJprr3bp/dYa/0pLjOR5Ei3z4yo8praH4rw6j0f+pJ9BHQiqaODGBA7GG2MH+r2v3a+hTVVMRNGJjv/iWC5VrlEwtm3XTKLDk9jzO5e6l3kvtO2eXL9z5caQvYLgbcS2ggKIWFgdE+mste8vgPyMsX6va/dq4ssojtIaaQlttCQlKEjQSB0AFaK7M3sa8RGP39YGHHsXmse/XHatsosTILLMEFlavGQm2pb02SPhDzDl6b69O+vBiarnw8yxm53PHbzdI92x22RWX4EJT7kR1hCg4w4kdW+YrCtnQ2Ds9OmdKVObRmKonExmfGZn3FA4CW2ZaOEGMw58R+BMajqDkaS2W3Gz2ijpST1B61f6rt+4dYtlE7x28Y7bLpL5A328uKhxfKO4bI3rqajveWwH5GWL9Xtfu1soprt0xRTETEbN/8C511PxVXFyJBb6uy5TMdA1vqpxI3+QbP5KjrFjNkw2E8zaLbCs0Ra+1cRFaSygq0BzHQA3oAb+asu8JcEffnsZJcWVMstJV4hHcSQslQ5S8oHu83YSPUpR9IrXpOkxotmble/o71p35ZfpSlfM1Kqcrh8nt3F2y93KxsrUVmLDDC2Qo9SUpdaXy7PXSSBsk661bKVsouVW/wyucKb5AXD5Z3v9BC/wAPTyAuHyzvf6CF/h6uVK3c5udnhH0Mqb5AXD5Z3v8AQQv8PTyAuHyzvf6CF/h6uVKc5udnhH0Mqgjh/IX5srKr1KZP2TX+TM8w9I52mUrH4UqB9RFWmHDYt0RmLFZRHjMoDbbTSQlKEgaAAHcK7qVrru13NlU+3oZyUpStKFKUoFKUoIXJMNs2XNIRdYKJC2wQ28CUOt77+VxJCk/kPWqU9wDtalks329R0HuQFsLA/AVNE/nJrJ9K7LWmaRYjVt1zELliz3gYPylvf5ov1FPeBg/KW9/mi/UVlOlb/iel/qen0MsWe8DB+Ut7/NF+op7wMH5S3v8ANF+orKdKfE9L/U9PoZYs94GD8pb3+aL9RXI4AwN9ckvZH/xR/wDhWUqU+J6X+p6GVKsHCDHLDIbkqYeuktshSHri52vKR3EI0EA/OEg1daUriu3rl6rWuVTM9pkpSlaUf//Z",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Image, display\n",
    "\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph import MessagesState\n",
    "from langgraph.graph import START, END, StateGraph\n",
    "from langgraph.prebuilt import tools_condition, ToolNode\n",
    "\n",
    "from langchain_core.messages import AIMessage, HumanMessage, SystemMessage\n",
    "\n",
    "# System message\n",
    "sys_msg = SystemMessage(content=\"You are a helpful assistant tasked with performing arithmetic on a set of inputs.\")\n",
    "\n",
    "# Node\n",
    "def assistant(state: MessagesState):\n",
    "   return {\"messages\": [llm_with_tools.invoke([sys_msg] + state[\"messages\"])]}\n",
    "\n",
    "# Graph\n",
    "builder = StateGraph(MessagesState)\n",
    "\n",
    "# Define nodes: these do the work\n",
    "builder.add_node(\"assistant\", assistant)\n",
    "builder.add_node(\"tools\", ToolNode(tools))\n",
    "\n",
    "# Define edges: these determine the control flow\n",
    "builder.add_edge(START, \"assistant\")\n",
    "builder.add_conditional_edges(\n",
    "    \"assistant\",\n",
    "    # If the latest message (result) from assistant is a tool call -> tools_condition routes to tools\n",
    "    # If the latest message (result) from assistant is a not a tool call -> tools_condition routes to END\n",
    "    tools_condition,\n",
    ")\n",
    "builder.add_edge(\"tools\", \"assistant\")\n",
    "\n",
    "memory = MemorySaver()\n",
    "graph = builder.compile(checkpointer=MemorySaver())\n",
    "\n",
    "# Show\n",
    "display(Image(graph.get_graph(xray=True).draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fab18a04-1329-47ac-a25b-4e01bf756e2a",
   "metadata": {},
   "source": [
    "Let's run it, as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "05b2ab62-82bc-4356-8d5b-2d4f49069fdd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "Multiply 2 and 3\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  multiply (call_ikJxMpb777bKMYgmM3d9mYjW)\n",
      " Call ID: call_ikJxMpb777bKMYgmM3d9mYjW\n",
      "  Args:\n",
      "    a: 2\n",
      "    b: 3\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: multiply\n",
      "\n",
      "6\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The result of multiplying 2 and 3 is 6.\n"
     ]
    }
   ],
   "source": [
    "# Input\n",
    "initial_input = {\"messages\": HumanMessage(content=\"Multiply 2 and 3\")}\n",
    "\n",
    "# Thread\n",
    "thread = {\"configurable\": {\"thread_id\": \"1\"}}\n",
    "\n",
    "# Run the graph until the first interruption\n",
    "for event in graph.stream(initial_input, thread, stream_mode=\"values\"):\n",
    "    event['messages'][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "268cfa43-22d1-4d63-8d81-a3ce00f1f2c8",
   "metadata": {},
   "source": [
    "## Browsing History\n",
    "\n",
    "We can use `get_state` to look at the **current** state of our graph, given the `thread_id`!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "161eb053-18f6-4c99-8674-8cbd11cae57e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "StateSnapshot(values={'messages': [HumanMessage(content='Multiply 2 and 3', id='4ee8c440-0e4a-47d7-852f-06e2a6c4f84d'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_ikJxMpb777bKMYgmM3d9mYjW', 'function': {'arguments': '{\"a\":2,\"b\":3}', 'name': 'multiply'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 17, 'prompt_tokens': 131, 'total_tokens': 148}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-bc24d334-8013-4f85-826f-e1ed69c86df0-0', tool_calls=[{'name': 'multiply', 'args': {'a': 2, 'b': 3}, 'id': 'call_ikJxMpb777bKMYgmM3d9mYjW', 'type': 'tool_call'}], usage_metadata={'input_tokens': 131, 'output_tokens': 17, 'total_tokens': 148}), ToolMessage(content='6', name='multiply', id='1012611a-30c5-4732-b789-8c455580c7b4', tool_call_id='call_ikJxMpb777bKMYgmM3d9mYjW'), AIMessage(content='The result of multiplying 2 and 3 is 6.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 156, 'total_tokens': 170}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5', 'finish_reason': 'stop', 'logprobs': None}, id='run-b46f3fed-ca3b-4e09-83f4-77ea5071e9bf-0', usage_metadata={'input_tokens': 156, 'output_tokens': 14, 'total_tokens': 170})]}, next=(), config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef6a440-ac9e-6024-8003-6fd8435c1d3b'}}, metadata={'source': 'loop', 'writes': {'assistant': {'messages': [AIMessage(content='The result of multiplying 2 and 3 is 6.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 156, 'total_tokens': 170}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5', 'finish_reason': 'stop', 'logprobs': None}, id='run-b46f3fed-ca3b-4e09-83f4-77ea5071e9bf-0', usage_metadata={'input_tokens': 156, 'output_tokens': 14, 'total_tokens': 170})]}}, 'step': 3, 'parents': {}}, created_at='2024-09-03T22:29:54.309727+00:00', parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef6a440-a759-6d02-8002-f1da6393e1ab'}}, tasks=())"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph.get_state({'configurable': {'thread_id': '1'}})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d00869e-7b41-4d71-ad3c-cacf8f9c029f",
   "metadata": {},
   "source": [
    "We can also browse the state history of our agent.\n",
    "\n",
    "`get_state_history` lets us get the state at all prior steps.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3010169c-3bfa-498c-a30c-7ba53744e4d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_states = [s for s in graph.get_state_history(thread)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c4612ccf-59fc-4848-8845-0433fee2ca8e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(all_states)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af30f269-1152-4fa1-a7c6-2947acad9a27",
   "metadata": {},
   "source": [
    "The first element is the current state, just as we got from `get_state`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4e60b292-8efc-4cc3-b836-51f060fa608b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "StateSnapshot(values={'messages': [HumanMessage(content='Multiply 2 and 3', id='4ee8c440-0e4a-47d7-852f-06e2a6c4f84d')]}, next=('assistant',), config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef6a440-a003-6c74-8000-8a2d82b0d126'}}, metadata={'source': 'loop', 'writes': None, 'step': 0, 'parents': {}}, created_at='2024-09-03T22:29:52.988265+00:00', parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef6a440-9ffe-6512-bfff-9e6d8dc24bba'}}, tasks=(PregelTask(id='ca669906-0c4f-5165-840d-7a6a3fce9fb9', name='assistant', error=None, interrupts=(), state=None),))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_states[-2]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4148a710-ceed-413b-b93c-070c6c792fa2",
   "metadata": {},
   "source": [
    "Everything above we can visualize here: \n",
    "\n",
    "![fig1.jpg](https://cdn.prod.website-files.com/65b8cd72835ceeacd4449a53/66dbb038211b544898570be3_time-travel1.png)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a5ad554a-faf3-489f-a9a9-774f4ec2a526",
   "metadata": {},
   "source": [
    "## Replaying \n",
    "\n",
    "We can re-run our agent from any of the prior steps.\n",
    "\n",
    "![fig2.jpg](https://cdn.prod.website-files.com/65b8cd72835ceeacd4449a53/66dbb038a0bd34b541c78fb8_time-travel2.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e135d2db-d613-42da-877e-d429f21aaefd",
   "metadata": {},
   "source": [
    "Let's look back at the step that recieved human input!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3688e511-a440-4330-a450-e5ed889c3b30",
   "metadata": {},
   "outputs": [],
   "source": [
    "to_replay = all_states[-2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "72adf296-d519-4bdc-af03-3b29799e9534",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "StateSnapshot(values={'messages': [HumanMessage(content='Multiply 2 and 3', id='4ee8c440-0e4a-47d7-852f-06e2a6c4f84d')]}, next=('assistant',), config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef6a440-a003-6c74-8000-8a2d82b0d126'}}, metadata={'source': 'loop', 'writes': None, 'step': 0, 'parents': {}}, created_at='2024-09-03T22:29:52.988265+00:00', parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef6a440-9ffe-6512-bfff-9e6d8dc24bba'}}, tasks=(PregelTask(id='ca669906-0c4f-5165-840d-7a6a3fce9fb9', name='assistant', error=None, interrupts=(), state=None),))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "to_replay"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "571e7894-6546-48ff-9c25-fa6d120391b3",
   "metadata": {},
   "source": [
    "Look at the state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6fe69428-f364-4330-bf5d-aa966c7f3b07",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content='Multiply 2 and 3', id='4ee8c440-0e4a-47d7-852f-06e2a6c4f84d')]}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "to_replay.values"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff2df545-cc80-4962-a34a-faac7af8eb3d",
   "metadata": {},
   "source": [
    "We can see the next node to call."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d2f333f9-9b2b-46f6-ac3a-525f86b20f1b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('assistant',)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "to_replay.next"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8938c18-5c22-47df-b71e-40afa73c87af",
   "metadata": {},
   "source": [
    "We also get the config, which tells us the `checkpoint_id` as well as the `thread_id`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b1298786-afa5-4277-927e-708a8629231b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'configurable': {'thread_id': '1',\n",
       "  'checkpoint_ns': '',\n",
       "  'checkpoint_id': '1ef6a440-a003-6c74-8000-8a2d82b0d126'}}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "to_replay.config"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d93b5eb-f541-4f82-93b1-48f54bf5cf83",
   "metadata": {},
   "source": [
    "To replay from here, we simply pass the config back to the agent!\n",
    "\n",
    "The graph knows that this checkpoint has aleady been executed. \n",
    "\n",
    "It just re-plays from this checkpoint!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "531b4cd1-54f6-44aa-9ffe-cf5403dad65d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "Multiply 2 and 3\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  multiply (call_SABfB57CnDkMu9HJeUE0mvJ9)\n",
      " Call ID: call_SABfB57CnDkMu9HJeUE0mvJ9\n",
      "  Args:\n",
      "    a: 2\n",
      "    b: 3\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: multiply\n",
      "\n",
      "6\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The result of multiplying 2 and 3 is 6.\n"
     ]
    }
   ],
   "source": [
    "for event in graph.stream(None, to_replay.config, stream_mode=\"values\"):\n",
    "    event['messages'][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d7a914e-63e6-4424-970f-15059ce9b4c3",
   "metadata": {},
   "source": [
    "Now, we can see our current state after the agent re-ran."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5a5a1f03-19f2-4d22-ba54-1c065ff08e85",
   "metadata": {},
   "source": [
    "## Forking\n",
    "\n",
    "What if we want to run from that same step, but with a different input.\n",
    "\n",
    "This is forking.\n",
    "\n",
    "![fig3.jpg](https://cdn.prod.website-files.com/65b8cd72835ceeacd4449a53/66dbb038f89f2d847ee5c336_time-travel3.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "cdeb5bf2-1566-4d8c-8ea5-65894e3a7038",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content='Multiply 2 and 3', id='4ee8c440-0e4a-47d7-852f-06e2a6c4f84d')]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "to_fork = all_states[-2]\n",
    "to_fork.values[\"messages\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a15f6a6-6eaa-48d6-92bb-864ea3a31b6a",
   "metadata": {},
   "source": [
    "Again, we have the config."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d1621b27-ee51-4dc3-81c4-1d05317280db",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'configurable': {'thread_id': '1',\n",
       "  'checkpoint_ns': '',\n",
       "  'checkpoint_id': '1ef6a440-a003-6c74-8000-8a2d82b0d126'}}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "to_fork.config"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2102195-0583-4dbe-ad2f-02fac7915585",
   "metadata": {},
   "source": [
    "Let's modify the state at this checkpoint.\n",
    "\n",
    "We can just run `update_state` with the `checkpoint_id` supplied. \n",
    "\n",
    "Remember how our reducer on `messages` works: \n",
    "\n",
    "* It will append, unless we supply a message ID.\n",
    "* We supply the message ID to overwrite the message, rather than appending to state!\n",
    "\n",
    "So, to overwrite the the message, we just supply the message ID, which we have `to_fork.values[\"messages\"].id`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0b4a918d-858a-41ac-a5d4-e99260e2d6ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "fork_config = graph.update_state(\n",
    "    to_fork.config,\n",
    "    {\"messages\": [HumanMessage(content='Multiply 5 and 3', \n",
    "                               id=to_fork.values[\"messages\"][0].id)]},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8ff4e9bb-8221-42d1-b7d0-b0cbd5dc374a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'configurable': {'thread_id': '1',\n",
       "  'checkpoint_ns': '',\n",
       "  'checkpoint_id': '1ef6a442-3661-62f6-8001-d3c01b96f98b'}}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fork_config"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bebfe6fd-c94b-4291-a125-ec6170e35bc5",
   "metadata": {},
   "source": [
    "This creates a new, forked checkpoint.\n",
    " \n",
    "But, the metadata - e.g., where to go next - is perserved! \n",
    "\n",
    "We can see the current state of our agent has been updated with our fork."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "586ce86c-1257-45e9-ba30-6287932b9484",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content='Multiply 5 and 3', id='4ee8c440-0e4a-47d7-852f-06e2a6c4f84d')]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_states = [state for state in graph.get_state_history(thread) ]\n",
    "all_states[0].values[\"messages\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "12e19798-25d8-49e8-8542-13d2b3bdf58e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "StateSnapshot(values={'messages': [HumanMessage(content='Multiply 5 and 3', id='4ee8c440-0e4a-47d7-852f-06e2a6c4f84d')]}, next=('assistant',), config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef6a442-3661-62f6-8001-d3c01b96f98b'}}, metadata={'source': 'update', 'step': 1, 'writes': {'__start__': {'messages': [HumanMessage(content='Multiply 5 and 3', id='4ee8c440-0e4a-47d7-852f-06e2a6c4f84d')]}}, 'parents': {}}, created_at='2024-09-03T22:30:35.598707+00:00', parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef6a440-a003-6c74-8000-8a2d82b0d126'}}, tasks=(PregelTask(id='f8990132-a8d3-5ddd-8d9e-1efbfc220da1', name='assistant', error=None, interrupts=(), state=None),))"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph.get_state({'configurable': {'thread_id': '1'}})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78c641e2-b8e9-4461-b854-8725006a5eb6",
   "metadata": {},
   "source": [
    "Now, when we stream, the graph knows this checkpoint has never been executed.\n",
    "\n",
    "So, the graph runs, rather than simply re-playing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "1c49f2a8-b325-45e4-b36c-17fab1b37cc0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "Multiply 5 and 3\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  multiply (call_KP2CVNMMUKMJAQuFmamHB21r)\n",
      " Call ID: call_KP2CVNMMUKMJAQuFmamHB21r\n",
      "  Args:\n",
      "    a: 5\n",
      "    b: 3\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: multiply\n",
      "\n",
      "15\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The result of multiplying 5 and 3 is 15.\n"
     ]
    }
   ],
   "source": [
    "for event in graph.stream(None, fork_config, stream_mode=\"values\"):\n",
    "    event['messages'][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "428d7f80-ee60-4147-b51f-ee3b0cf5cbba",
   "metadata": {},
   "source": [
    "Now, we can see the current state is the end of our agent run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "132ef840-64c7-479c-ad34-3f177f4b2524",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "StateSnapshot(values={'messages': [HumanMessage(content='Multiply 5 and 3', id='4ee8c440-0e4a-47d7-852f-06e2a6c4f84d'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_KP2CVNMMUKMJAQuFmamHB21r', 'function': {'arguments': '{\"a\":5,\"b\":3}', 'name': 'multiply'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 17, 'prompt_tokens': 131, 'total_tokens': 148}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-bc420009-d1f6-49b8-bea7-dfc9fca7eb79-0', tool_calls=[{'name': 'multiply', 'args': {'a': 5, 'b': 3}, 'id': 'call_KP2CVNMMUKMJAQuFmamHB21r', 'type': 'tool_call'}], usage_metadata={'input_tokens': 131, 'output_tokens': 17, 'total_tokens': 148}), ToolMessage(content='15', name='multiply', id='9232e653-816d-471a-9002-9a1ecd453364', tool_call_id='call_KP2CVNMMUKMJAQuFmamHB21r'), AIMessage(content='The result of multiplying 5 and 3 is 15.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 156, 'total_tokens': 170}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5', 'finish_reason': 'stop', 'logprobs': None}, id='run-86c21888-d832-47c5-9e76-0aa2676116dc-0', usage_metadata={'input_tokens': 156, 'output_tokens': 14, 'total_tokens': 170})]}, next=(), config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef6a442-a2e2-6e98-8004-4a0b75537950'}}, metadata={'source': 'loop', 'writes': {'assistant': {'messages': [AIMessage(content='The result of multiplying 5 and 3 is 15.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 156, 'total_tokens': 170}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5', 'finish_reason': 'stop', 'logprobs': None}, id='run-86c21888-d832-47c5-9e76-0aa2676116dc-0', usage_metadata={'input_tokens': 156, 'output_tokens': 14, 'total_tokens': 170})]}}, 'step': 4, 'parents': {}}, created_at='2024-09-03T22:30:46.976463+00:00', parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef6a442-9db0-6056-8003-7304cab7bed8'}}, tasks=())"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph.get_state({'configurable': {'thread_id': '1'}})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2ceb5f31-97b0-466c-9b3b-ae4df7cd462a",
   "metadata": {},
   "source": [
    "### Time travel with LangGraph API\n",
    "\n",
    "**⚠️ DISCLAIMER**\n",
    "\n",
    "Since the filming of these videos, we've updated Studio so that it can be run locally and opened in your browser. This is now the preferred way to run Studio (rather than using the Desktop App as shown in the video). See documentation [here](https://langchain-ai.github.io/langgraph/concepts/langgraph_studio/#local-development-server) on the local development server and [here](https://langchain-ai.github.io/langgraph/how-tos/local-studio/#run-the-development-server). To start the local development server, run the following command in your terminal in the `/studio` directory in this module:\n",
    "\n",
    "```\n",
    "langgraph dev\n",
    "```\n",
    "\n",
    "You should see the following output:\n",
    "```\n",
    "- 🚀 API: http://127.0.0.1:2024\n",
    "- 🎨 Studio UI: https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024\n",
    "- 📚 API Docs: http://127.0.0.1:2024/docs\n",
    "```\n",
    "\n",
    "Open your browser and navigate to the Studio UI: `https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024`.\n",
    "\n",
    "We connect to it via the SDK and show how the LangGraph API [supports time travel](https://langchain-ai.github.io/langgraph/cloud/how-tos/human_in_the_loop_time_travel/#initial-invocation). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "891defdb-746c-48e4-8efa-bb5f138dc4bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "if 'google.colab' in str(get_ipython()):\n",
    "    raise Exception(\"Unfortunately LangGraph Studio is currently not supported on Google Colab\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "a317925d-1788-4cfc-9c12-336b17b4d859",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph_sdk import get_client\n",
    "client = get_client(url=\"http://127.0.0.1:2024\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "815d5e03-0ab8-4c7f-a1ee-f410b6aadc03",
   "metadata": {},
   "source": [
    "#### Re-playing \n",
    "\n",
    "Let's run our agent streaming `updates` to the state of the graph after each node is called."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "9d4d01da-7b64-4c92-96b7-29ec93332d0b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------Assistant Node--------------------\n",
      "{'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': 'call_SG7XYqDENGq7mwXrnioNLosS', 'function': {'arguments': '{\"a\":2,\"b\":3}', 'name': 'multiply'}, 'type': 'function'}]}, 'response_metadata': {'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-2c120fc3-3c82-4599-b8ec-24fbee207cad', 'example': False, 'tool_calls': [{'name': 'multiply', 'args': {'a': 2, 'b': 3}, 'id': 'call_SG7XYqDENGq7mwXrnioNLosS', 'type': 'tool_call'}], 'invalid_tool_calls': [], 'usage_metadata': None}\n",
      "--------------------Tools Node--------------------\n",
      "{'content': '6', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'tool', 'name': 'multiply', 'id': '3b40d091-58b2-4566-a84c-60af67206307', 'tool_call_id': 'call_SG7XYqDENGq7mwXrnioNLosS', 'artifact': None, 'status': 'success'}\n",
      "--------------------Assistant Node--------------------\n",
      "{'content': 'The result of multiplying 2 and 3 is 6.', 'additional_kwargs': {}, 'response_metadata': {'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_fde2829a40'}, 'type': 'ai', 'name': None, 'id': 'run-1272d9b0-a0aa-4ff7-8bad-fdffd27c5506', 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None}\n"
     ]
    }
   ],
   "source": [
    "initial_input = {\"messages\": HumanMessage(content=\"Multiply 2 and 3\")}\n",
    "thread = await client.threads.create()\n",
    "async for chunk in client.runs.stream(\n",
    "    thread[\"thread_id\"],\n",
    "    assistant_id = \"agent\",\n",
    "    input=initial_input,\n",
    "    stream_mode=\"updates\",\n",
    "):\n",
    "    if chunk.data:\n",
    "        assisant_node = chunk.data.get('assistant', {}).get('messages', [])\n",
    "        tool_node = chunk.data.get('tools', {}).get('messages', [])\n",
    "        if assisant_node:\n",
    "            print(\"-\" * 20+\"Assistant Node\"+\"-\" * 20)\n",
    "            print(assisant_node[-1])\n",
    "        elif tool_node:\n",
    "            print(\"-\" * 20+\"Tools Node\"+\"-\" * 20)\n",
    "            print(tool_node[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cc3bab2",
   "metadata": {},
   "source": [
    "Now, let's look at **replaying** from a specified checkpoint. \n",
    "\n",
    "We simply need to pass the `checkpoint_id`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "d8ecc4fd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'values': {'messages': [{'content': 'Multiply 2 and 3',\n",
       "    'additional_kwargs': {'example': False,\n",
       "     'additional_kwargs': {},\n",
       "     'response_metadata': {}},\n",
       "    'response_metadata': {},\n",
       "    'type': 'human',\n",
       "    'name': None,\n",
       "    'id': 'df98147a-cb3d-4f1a-b7f7-1545c4b6f042',\n",
       "    'example': False}]},\n",
       " 'next': ['assistant'],\n",
       " 'tasks': [{'id': 'e497456f-827a-5027-87bd-b0ccd54aa89a',\n",
       "   'name': 'assistant',\n",
       "   'error': None,\n",
       "   'interrupts': [],\n",
       "   'state': None}],\n",
       " 'metadata': {'step': 0,\n",
       "  'run_id': '1ef6a449-7fbc-6c90-8754-4e6b1b582790',\n",
       "  'source': 'loop',\n",
       "  'writes': None,\n",
       "  'parents': {},\n",
       "  'user_id': '',\n",
       "  'graph_id': 'agent',\n",
       "  'thread_id': '708e1d8f-f7c8-4093-9bb4-999c4237cb4a',\n",
       "  'created_by': 'system',\n",
       "  'assistant_id': 'fe096781-5601-53d2-b2f6-0d3403f7e9ca'},\n",
       " 'created_at': '2024-09-03T22:33:51.380352+00:00',\n",
       " 'checkpoint_id': '1ef6a449-817f-6b55-8000-07c18fbdf7c8',\n",
       " 'parent_checkpoint_id': '1ef6a449-816c-6fd6-bfff-32a56dd2635f'}"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "states = await client.threads.get_history(thread['thread_id'])\n",
    "to_replay = states[-2]\n",
    "to_replay"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e33f865a",
   "metadata": {},
   "source": [
    "Let's stream with `stream_mode=\"values\"` to see the full state at every node as we replay. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "325e8272",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Receiving new event of type: metadata...\n",
      "{'run_id': '1ef6a44a-5806-6bb1-b2ee-92ecfda7f67d'}\n",
      "\n",
      "\n",
      "\n",
      "Receiving new event of type: values...\n",
      "{'messages': [{'content': 'Multiply 2 and 3', 'additional_kwargs': {'example': False, 'additional_kwargs': {}, 'response_metadata': {}}, 'response_metadata': {}, 'type': 'human', 'name': None, 'id': 'df98147a-cb3d-4f1a-b7f7-1545c4b6f042', 'example': False}]}\n",
      "\n",
      "\n",
      "\n",
      "Receiving new event of type: values...\n",
      "{'messages': [{'content': 'Multiply 2 and 3', 'additional_kwargs': {'example': False, 'additional_kwargs': {}, 'response_metadata': {}}, 'response_metadata': {}, 'type': 'human', 'name': None, 'id': 'df98147a-cb3d-4f1a-b7f7-1545c4b6f042', 'example': False}, {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': 'call_Rn9YQ6iZyYtzrELBz7EfQcs0', 'function': {'arguments': '{\"a\":2,\"b\":3}', 'name': 'multiply'}, 'type': 'function'}]}, 'response_metadata': {'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-e60d82d7-7743-4f13-bebd-3616a88720a9', 'example': False, 'tool_calls': [{'name': 'multiply', 'args': {'a': 2, 'b': 3}, 'id': 'call_Rn9YQ6iZyYtzrELBz7EfQcs0', 'type': 'tool_call'}], 'invalid_tool_calls': [], 'usage_metadata': None}]}\n",
      "\n",
      "\n",
      "\n",
      "Receiving new event of type: values...\n",
      "{'messages': [{'content': 'Multiply 2 and 3', 'additional_kwargs': {'example': False, 'additional_kwargs': {}, 'response_metadata': {}}, 'response_metadata': {}, 'type': 'human', 'name': None, 'id': 'df98147a-cb3d-4f1a-b7f7-1545c4b6f042', 'example': False}, {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': 'call_Rn9YQ6iZyYtzrELBz7EfQcs0', 'function': {'arguments': '{\"a\":2,\"b\":3}', 'name': 'multiply'}, 'type': 'function'}]}, 'response_metadata': {'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-e60d82d7-7743-4f13-bebd-3616a88720a9', 'example': False, 'tool_calls': [{'name': 'multiply', 'args': {'a': 2, 'b': 3}, 'id': 'call_Rn9YQ6iZyYtzrELBz7EfQcs0', 'type': 'tool_call'}], 'invalid_tool_calls': [], 'usage_metadata': None}, {'content': '6', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'tool', 'name': 'multiply', 'id': 'f1be0b83-4565-4aa2-9b9a-cd8874c6a2bc', 'tool_call_id': 'call_Rn9YQ6iZyYtzrELBz7EfQcs0', 'artifact': None, 'status': 'success'}]}\n",
      "\n",
      "\n",
      "\n",
      "Receiving new event of type: values...\n",
      "{'messages': [{'content': 'Multiply 2 and 3', 'additional_kwargs': {'example': False, 'additional_kwargs': {}, 'response_metadata': {}}, 'response_metadata': {}, 'type': 'human', 'name': None, 'id': 'df98147a-cb3d-4f1a-b7f7-1545c4b6f042', 'example': False}, {'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': 'call_Rn9YQ6iZyYtzrELBz7EfQcs0', 'function': {'arguments': '{\"a\":2,\"b\":3}', 'name': 'multiply'}, 'type': 'function'}]}, 'response_metadata': {'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-e60d82d7-7743-4f13-bebd-3616a88720a9', 'example': False, 'tool_calls': [{'name': 'multiply', 'args': {'a': 2, 'b': 3}, 'id': 'call_Rn9YQ6iZyYtzrELBz7EfQcs0', 'type': 'tool_call'}], 'invalid_tool_calls': [], 'usage_metadata': None}, {'content': '6', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'tool', 'name': 'multiply', 'id': 'f1be0b83-4565-4aa2-9b9a-cd8874c6a2bc', 'tool_call_id': 'call_Rn9YQ6iZyYtzrELBz7EfQcs0', 'artifact': None, 'status': 'success'}, {'content': 'The result of multiplying 2 and 3 is 6.', 'additional_kwargs': {}, 'response_metadata': {'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-55e5847a-d542-4977-84d7-24852e78b0a9', 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None}]}\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "async for chunk in client.runs.stream(\n",
    "    thread[\"thread_id\"],\n",
    "    assistant_id=\"agent\",\n",
    "    input=None,\n",
    "    stream_mode=\"values\",\n",
    "    checkpoint_id=to_replay['checkpoint_id']\n",
    "):      \n",
    "    print(f\"Receiving new event of type: {chunk.event}...\")\n",
    "    print(chunk.data)\n",
    "    print(\"\\n\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14c153b3",
   "metadata": {},
   "source": [
    "We can all view this as streaming only `updates` to state made by the nodes that we reply."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "9e608e93",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------Assistant Node--------------------\n",
      "{'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': 'call_I2qudhMCwcw1GzcFN5q80rjj', 'function': {'arguments': '{\"a\":2,\"b\":3}', 'name': 'multiply'}, 'type': 'function'}]}, 'response_metadata': {'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-550e75ad-dbbc-4e55-9f00-aa896228914c', 'example': False, 'tool_calls': [{'name': 'multiply', 'args': {'a': 2, 'b': 3}, 'id': 'call_I2qudhMCwcw1GzcFN5q80rjj', 'type': 'tool_call'}], 'invalid_tool_calls': [], 'usage_metadata': None}\n",
      "--------------------Tools Node--------------------\n",
      "{'content': '6', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'tool', 'name': 'multiply', 'id': '731b7d4f-780d-4a8b-aec9-0d8b9c58c40a', 'tool_call_id': 'call_I2qudhMCwcw1GzcFN5q80rjj', 'artifact': None, 'status': 'success'}\n",
      "--------------------Assistant Node--------------------\n",
      "{'content': 'The result of multiplying 2 and 3 is 6.', 'additional_kwargs': {}, 'response_metadata': {'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-2326afa5-eb43-4568-b5ed-424c0a0fa076', 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None}\n"
     ]
    }
   ],
   "source": [
    "async for chunk in client.runs.stream(\n",
    "    thread[\"thread_id\"],\n",
    "    assistant_id=\"agent\",\n",
    "    input=None,\n",
    "    stream_mode=\"updates\",\n",
    "    checkpoint_id=to_replay['checkpoint_id']\n",
    "):\n",
    "    if chunk.data:\n",
    "        assisant_node = chunk.data.get('assistant', {}).get('messages', [])\n",
    "        tool_node = chunk.data.get('tools', {}).get('messages', [])\n",
    "        if assisant_node:\n",
    "            print(\"-\" * 20+\"Assistant Node\"+\"-\" * 20)\n",
    "            print(assisant_node[-1])\n",
    "        elif tool_node:\n",
    "            print(\"-\" * 20+\"Tools Node\"+\"-\" * 20)\n",
    "            print(tool_node[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e66e0e8",
   "metadata": {},
   "source": [
    "#### Forking\n",
    "\n",
    "Now, let's look at forking.\n",
    "\n",
    "Let's get the same step as we worked with above, the human input.\n",
    "\n",
    "Let's create a new thread with our agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "01af5ed4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------Assistant Node--------------------\n",
      "{'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': 'call_HdWoyLELFZGEcqGxFt2fZzek', 'function': {'arguments': '{\"a\":2,\"b\":3}', 'name': 'multiply'}, 'type': 'function'}]}, 'response_metadata': {'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-cbd081b1-8cef-4ca8-9dd5-aceb134404dc', 'example': False, 'tool_calls': [{'name': 'multiply', 'args': {'a': 2, 'b': 3}, 'id': 'call_HdWoyLELFZGEcqGxFt2fZzek', 'type': 'tool_call'}], 'invalid_tool_calls': [], 'usage_metadata': None}\n",
      "--------------------Tools Node--------------------\n",
      "{'content': '6', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'tool', 'name': 'multiply', 'id': '11dd4a7f-0b6b-44da-b9a4-65f1677c8813', 'tool_call_id': 'call_HdWoyLELFZGEcqGxFt2fZzek', 'artifact': None, 'status': 'success'}\n",
      "--------------------Assistant Node--------------------\n",
      "{'content': 'The result of multiplying 2 and 3 is 6.', 'additional_kwargs': {}, 'response_metadata': {'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-936cf990-9302-45c7-9051-6ff1e2e9f316', 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None}\n"
     ]
    }
   ],
   "source": [
    "initial_input = {\"messages\": HumanMessage(content=\"Multiply 2 and 3\")}\n",
    "thread = await client.threads.create()\n",
    "async for chunk in client.runs.stream(\n",
    "    thread[\"thread_id\"],\n",
    "    assistant_id=\"agent\",\n",
    "    input=initial_input,\n",
    "    stream_mode=\"updates\",\n",
    "):\n",
    "    if chunk.data:\n",
    "        assisant_node = chunk.data.get('assistant', {}).get('messages', [])\n",
    "        tool_node = chunk.data.get('tools', {}).get('messages', [])\n",
    "        if assisant_node:\n",
    "            print(\"-\" * 20+\"Assistant Node\"+\"-\" * 20)\n",
    "            print(assisant_node[-1])\n",
    "        elif tool_node:\n",
    "            print(\"-\" * 20+\"Tools Node\"+\"-\" * 20)\n",
    "            print(tool_node[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "3dbc8795-c3f5-4559-a00e-dc410c0a927f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [{'content': 'Multiply 2 and 3',\n",
       "   'additional_kwargs': {'example': False,\n",
       "    'additional_kwargs': {},\n",
       "    'response_metadata': {}},\n",
       "   'response_metadata': {},\n",
       "   'type': 'human',\n",
       "   'name': None,\n",
       "   'id': '93c18b95-9050-4a52-99b8-9374e98ee5db',\n",
       "   'example': False}]}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "states = await client.threads.get_history(thread['thread_id'])\n",
    "to_fork = states[-2]\n",
    "to_fork['values']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "11e6cde1-0388-43ea-b994-1c4e9ca1199b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'93c18b95-9050-4a52-99b8-9374e98ee5db'"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "to_fork['values']['messages'][0]['id']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "0c1e2300-c8b2-4994-a96d-1be19c04b6a8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['assistant']"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "to_fork['next']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "9d31d5aa-524f-42f4-ba7e-713a029610d6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1ef6a44b-27ec-681c-8000-ff7e345aee7e'"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "to_fork['checkpoint_id']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f11e1d9-9fe7-4243-a06f-9b07e38a12ad",
   "metadata": {},
   "source": [
    "Let's edit the state.\n",
    "\n",
    "Remember how our reducer on `messages` works: \n",
    "\n",
    "* It will append, unless we supply a message ID.\n",
    "* We supply the message ID to overwrite the message, rather than appending to state!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "0198f1b8-2f57-4c6e-ac6a-c6fb80cce0bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "forked_input = {\"messages\": HumanMessage(content=\"Multiply 3 and 3\",\n",
    "                                         id=to_fork['values']['messages'][0]['id'])}\n",
    "\n",
    "forked_config = await client.threads.update_state(\n",
    "    thread[\"thread_id\"],\n",
    "    forked_input,\n",
    "    checkpoint_id=to_fork['checkpoint_id']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "1dcd5b8e-6bb1-4967-84cf-4af710b8bf46",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'configurable': {'thread_id': 'c99502e7-b0d7-473e-8295-1ad60e2b7ed2',\n",
       "  'checkpoint_ns': '',\n",
       "  'checkpoint_id': '1ef6a44b-90dc-68c8-8001-0c36898e0f34'},\n",
       " 'checkpoint_id': '1ef6a44b-90dc-68c8-8001-0c36898e0f34'}"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "forked_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "015ac68a-5cc1-4c42-90a2-5b2b4865a153",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'values': {'messages': [{'content': 'Multiply 3 and 3',\n",
       "    'additional_kwargs': {'additional_kwargs': {},\n",
       "     'response_metadata': {},\n",
       "     'example': False},\n",
       "    'response_metadata': {},\n",
       "    'type': 'human',\n",
       "    'name': None,\n",
       "    'id': '93c18b95-9050-4a52-99b8-9374e98ee5db',\n",
       "    'example': False}]},\n",
       " 'next': ['assistant'],\n",
       " 'tasks': [{'id': 'da5d6548-62ca-5e69-ba70-f6179b2743bd',\n",
       "   'name': 'assistant',\n",
       "   'error': None,\n",
       "   'interrupts': [],\n",
       "   'state': None}],\n",
       " 'metadata': {'step': 1,\n",
       "  'source': 'update',\n",
       "  'writes': {'__start__': {'messages': {'id': '93c18b95-9050-4a52-99b8-9374e98ee5db',\n",
       "     'name': None,\n",
       "     'type': 'human',\n",
       "     'content': 'Multiply 3 and 3',\n",
       "     'example': False,\n",
       "     'additional_kwargs': {},\n",
       "     'response_metadata': {}}}},\n",
       "  'parents': {},\n",
       "  'graph_id': 'agent'},\n",
       " 'created_at': '2024-09-03T22:34:46.678333+00:00',\n",
       " 'checkpoint_id': '1ef6a44b-90dc-68c8-8001-0c36898e0f34',\n",
       " 'parent_checkpoint_id': '1ef6a44b-27ec-681c-8000-ff7e345aee7e'}"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "states = await client.threads.get_history(thread['thread_id'])\n",
    "states[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3de80029-b987-49c5-890d-6cd70cbc8de7",
   "metadata": {},
   "source": [
    "To rerun, we pass in the `checkpoint_id`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "da005240-d3f0-4c89-9aca-dfcb5d410ceb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------Assistant Node--------------------\n",
      "{'content': '', 'additional_kwargs': {'tool_calls': [{'index': 0, 'id': 'call_aodhCt5fWv33qVbO7Nsub9Q3', 'function': {'arguments': '{\"a\":3,\"b\":3}', 'name': 'multiply'}, 'type': 'function'}]}, 'response_metadata': {'finish_reason': 'tool_calls', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-e9759422-e537-4b9b-b583-36c688e13b4b', 'example': False, 'tool_calls': [{'name': 'multiply', 'args': {'a': 3, 'b': 3}, 'id': 'call_aodhCt5fWv33qVbO7Nsub9Q3', 'type': 'tool_call'}], 'invalid_tool_calls': [], 'usage_metadata': None}\n",
      "--------------------Tools Node--------------------\n",
      "{'content': '9', 'additional_kwargs': {}, 'response_metadata': {}, 'type': 'tool', 'name': 'multiply', 'id': '89787b0b-93de-4c0a-bea8-d2c3845534e1', 'tool_call_id': 'call_aodhCt5fWv33qVbO7Nsub9Q3', 'artifact': None, 'status': 'success'}\n",
      "--------------------Assistant Node--------------------\n",
      "{'content': 'The result of multiplying 3 by 3 is 9.', 'additional_kwargs': {}, 'response_metadata': {'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_157b3831f5'}, 'type': 'ai', 'name': None, 'id': 'run-0e16610f-4e8d-46f3-a5df-c2f187fae593', 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None}\n"
     ]
    }
   ],
   "source": [
    "async for chunk in client.runs.stream(\n",
    "    thread[\"thread_id\"],\n",
    "    assistant_id=\"agent\",\n",
    "    input=None,\n",
    "    stream_mode=\"updates\",\n",
    "    checkpoint_id=forked_config['checkpoint_id']\n",
    "):\n",
    "    if chunk.data:\n",
    "        assisant_node = chunk.data.get('assistant', {}).get('messages', [])\n",
    "        tool_node = chunk.data.get('tools', {}).get('messages', [])\n",
    "        if assisant_node:\n",
    "            print(\"-\" * 20+\"Assistant Node\"+\"-\" * 20)\n",
    "            print(assisant_node[-1])\n",
    "        elif tool_node:\n",
    "            print(\"-\" * 20+\"Tools Node\"+\"-\" * 20)\n",
    "            print(tool_node[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36956571-a2b8-4f1b-8e30-51f02f155a6f",
   "metadata": {},
   "source": [
    "### LangGraph Studio\n",
    "\n",
    "Let's look at forking in the Studio UI with our `agent`, which uses `module-1/studio/agent.py` set in `module-1/studio/langgraph.json`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
